# -*- coding: utf-8 -*-
"""
@author: YuHaiyang

"""

import torch
from PIL import Image
from torchvision import transforms
from torchvision.models import VGG11_Weights

# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    # net: torch.nn.Module = torch.hub.load("pytorch/vision:v0.10.0", 'vgg11', pretrained=True)
    net: torch.nn.Module = torch.hub.load("pytorch/vision:v0.10.0", 'vgg11', weights=VGG11_Weights.DEFAULT)
    net.eval()

    image = Image.open("../../assets/dog.jpg")

    transform = transforms.Compose(
        [
            transforms.Resize(225),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ]
    )

    image = transform(image)
    image = image.unsqueeze(0)

    if torch.cuda.is_available():
        image = image.to('cuda')
        net.to('cuda')
    elif torch.backends.mps.is_available():
        image = image.to('mps')
        net.to('mps')

    with torch.no_grad():
        out = net(image)

    probabilities = torch.nn.functional.softmax(out[0], dim=0)

    with open("../../assets/imagenet_classes.txt", "r") as f:
        categories = [s.strip() for s in f.readlines()]
    # Show top categories per image
    top5_prob, top5_catid = torch.topk(probabilities, 5)
    for i in range(top5_prob.size(0)):
        print(categories[top5_catid[i]], top5_prob[i].item())
